import os
import shutil
import numpy as np
from astropy.io import fits
from scipy.signal import savgol_filter
from scipy.stats import trim_mean, sigmaclip
from tqdm import trange


def _smooth_data(data, filterwidth):
    dd = savgol_filter(data, filterwidth * 4 + 1, 3)
    diff = data - dd
    index = np.abs(diff) > sigmaclip(diff, 5, 5)[0].std() * 5
    x = np.arange(len(data))
    dd = np.interp(x, x[~index], data[~index])
    return savgol_filter(dd, filterwidth * 2 + 1, 2)


def extract_overscan(image, os_x1, os_x2, y1, y2, filterwidth=5):
    if os_x1 >= os_x2:
        raise Exception('os_x1 must be smaller than os_x2')
    if y1 >= y2:
        raise Exception('y1 must be smaller than y2')
    data = trim_mean(image[y1:y2, os_x1:os_x2], 0.1, axis=1)
    if filterwidth > 1:
        data = _smooth_data(data, filterwidth)
    return data


def subtract_overscan(image, x0, nx, y0, ny, os_x1, os_x2):
    v = extract_overscan(image, os_x1, os_x2, y0, y0 + ny)
    v = np.repeat(v, nx).reshape((ny, nx))
    image[y0:y0 + ny, x0:x0 + nx] -= v
    return image


def corr_overscan(image, imgconf):
    for chan in range(imgconf.nchan):
        x0 = imgconf.raw_x0[chan]
        y0 = imgconf.raw_y0[chan]
        if imgconf.raw_xflip[chan]:
            x1 = x0 + imgconf.raw_nx[chan] - imgconf.x1[chan] - imgconf.nx[chan]
            os1 = x0 + imgconf.raw_nx[chan] - imgconf.overscan_x2[chan]
            os2 = x0 + imgconf.raw_nx[chan] - imgconf.overscan_x1[chan]
        else:
            x1 = x0 + imgconf.x1[chan]
            os1 = x0 + imgconf.overscan_x1[chan]
            os2 = x0 + imgconf.overscan_x2[chan]
        if imgconf.raw_yflip[chan]:
            y1 = y0 + imgconf.raw_ny[chan] - imgconf.y1[chan] - imgconf.ny[chan]
        else:
            y1 = y0 + imgconf.y1[chan]
        image = subtract_overscan(image, x1, imgconf.nx[chan], y1, imgconf.ny[chan], os1, os2)
    return image


def crop_overscan(image, imgconf):
    output = np.zeros(imgconf.mos_shape, dtype=image.dtype)
    for chan in range(imgconf.nchan):
        x1 = imgconf.raw_x0[chan]
        y1 = imgconf.raw_y0[chan]
        cimage = image[y1:y1 + imgconf.raw_ny[chan], x1:x1 + imgconf.raw_nx[chan]]

        if imgconf.raw_xflip[chan]:
            x1 = imgconf.raw_nx[chan] - imgconf.x1[chan] - imgconf.nx[chan]
        else:
            x1 = imgconf.x1[chan]
        if imgconf.raw_yflip[chan]:
            y1 = imgconf.raw_ny[chan] - imgconf.y1[chan] - imgconf.ny[chan]
        else:
            y1 = imgconf.y1[chan]
        cimage = cimage[y1:y1 + imgconf.ny[chan], x1:x1 + imgconf.nx[chan]]

        if imgconf.mos_xflip[chan]:
            cimage = cimage[:, ::-1]
        if imgconf.mos_yflip[chan]:
            cimage = cimage[::-1, :]
        x1 = imgconf.mos_x0[chan]
        y1 = imgconf.mos_y0[chan]
        output[y1:y1 + imgconf.ny[chan], x1:x1 + imgconf.nx[chan]] = cimage
    return output


def do_overscan(flist, imgconf, crop=True, outlist=None, croponly=False,
                outdir=None, keepname=True, outname='tmp', copyfile=False, flat_flag=False, flat_file=None):
    print(imgconf)
    # if not necessary, just do nothing
    if imgconf.overscan == 0 and copyfile is False:
        print('\n############################\n')
        print('## NOT doing overscan ######')
        # print('## return flist: \n{}\n'.format(flist))
        print('############################\n')
        return flist
    if croponly:
        crop = True


    # file name
    n = len(flist)
    if outlist is None:
        if outdir is None:
            raise Exception('Please define output directory or output file list')
        outlist = []
        for i in range(n):
            if keepname is True:
                outlist.append('{}/{}'.format(outdir, os.path.basename(flist[i])))
            else:
                ndigit = int(np.log10(n)) + 1
                txt = '{}/{}-s{:0' + str(ndigit) + 'd}.fits'
                outlist.append(txt.format(outdir, outname, i + 1))
    if len(outlist) != n:
        raise Exception('output name list is not as long as input list')

    # do overscan correction
    for i in trange(n, desc='image cropping/overscan correction'):
        if imgconf.overscan == 1:
            with fits.open(flist[i]) as img:
                data = img[0].data.astype(np.float32)
                if croponly is False and imgconf.type != 'cmos':
                    data = corr_overscan(data, imgconf)
                if crop:
                    data = crop_overscan(data, imgconf)
                hdu = fits.PrimaryHDU(data, header=img[0].header)
                hdu.header['ORIGNAME'] = flist[i]
                hdu.writeto(outlist[i], overwrite=True)
        else:
            shutil.copy(flist[i], outlist[i])

    return outlist


def fill_overscan(image, imgconf, fill_value=0):
    output = np.full(imgconf.raw_shape, fill_value, dtype=image.dtype)
    for chan in range(imgconf.nchan):
        x1 = imgconf.mos_x0[chan]
        y1 = imgconf.mos_y0[chan]
        cimage = image[y1:y1 + imgconf.ny[chan], x1:x1 + imgconf.nx[chan]]
        if imgconf.mos_xflip[chan]:
            cimage = cimage[:, ::-1]
        if imgconf.mos_yflip[chan]:
            cimage = cimage[::-1, :]

        if imgconf.raw_xflip[chan]:
            x2 = imgconf.raw_x0[chan] + imgconf.raw_nx[chan] - imgconf.x1[chan]
            x1 = x2 - imgconf.nx[chan]
        else:
            x1 = imgconf.raw_x0[chan] + imgconf.x1[chan]
            x2 = x1 + imgconf.nx[chan]
        if imgconf.raw_yflip[chan]:
            y2 = imgconf.raw_y0[chan] + imgconf.raw_ny[chan] - imgconf.y1[chan]
            y1 = y2 - imgconf.ny[chan]
        else:
            y1 = imgconf.raw_y0[chan] + imgconf.y1[chan]
            y2 = y1 + imgconf.ny[chan]
        output[y1:y2, x1:x2] = cimage

    return output
